#pragma once
#include <lcir/ADT/GraphNode.h>
#include <lcir/IR/Expression.h>
#include <lcir/IR/IRRewritter.h>
#include <lcir/IR/Op.h>
#include <lcir/Utils/Result.h>

#include <cstddef>
#include <map>
#include <vector>
namespace lcir {
class PatternNode {
  std::vector<PatternNode> inputs;
  OpInstance *opInstance{nullptr};

 public:
  PatternNode() = default;
  PatternNode(OpInstance *opInstance) : opInstance(opInstance) {}
  Result match(Expression *expr) const;
  std::vector<PatternNode> getInputs() const { return inputs; };
  void addInput(PatternNode node) { inputs.push_back(node); }
  void addInput(OpInstance *instance) { inputs.push_back(instance); }
  void setInputs(const std::vector<PatternNode> &inputs) { this->inputs = inputs; }
  void setInputsWithInstance(const std::vector<OpInstance *> &instances) {
    for (auto *i : instances) inputs.push_back(i);
  }
  OpInstance *getOp() const { return opInstance; }
};

class PatternBase {
 protected:
  size_t cost = -1;

 public:
  virtual Result match(Expression *root) const = 0;
  virtual Result rewrite(Expression *root, IRRewritter &b) const = 0;
  virtual Result matchAndRewrite(Expression *root) const = 0;
  size_t getCost() const { return cost; }
  void setCost(size_t cost) { this->cost = cost; }
};

class Pattern : public PatternBase {
  PatternNode root;
 public:
  Pattern() = default;
  void setRootNode(PatternNode root) { this->root = root; }
  OpInstance *getRootOp() {return root.getOp();}
  
  /**
   * @brief 匹配表达式
   * @param Expression表达式根节点
   */
  Result match(Expression *root) const override;
  /**
   * @brief 匹配并重写表达式
   * @param Expression表达式根节点
   * @warning 不要在这个方法内erase除root外的其他节点！
   */
  Result matchAndRewrite(Expression *root) const override;

  /**
   * @param
   */
  virtual Result rewrite(Expression *root, IRRewritter &b) const override { return failure(); }
};

/**
 * @brief 匹配单个Op节点的Pattern
 */
template<class Op>
class OpPattern : public Pattern {
  OpInstance *op;
 public:
  OpPattern() { setRootNode(PatternNode(Op::getOpInstance())); }
};

/**
 * TODO:这里面能同时放多个pattern
 */
// class MultiPattern : public PatternBase {
//   std::map<OpInstance *, PatternNode> opToRoot;

//  public:
//   MultiPattern(PatternNode root, size_t cost = -1) : PatternBase(cost) {}
//   Result match(Expression *root) const override { return success(); }
//   Resutt matchAndResrite()
//   virtual Result rewrite(Expression *root, IRRewritter &b) const override { return failure(); }
// };
}  // namespace lcir